import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter

def ssr_channel_raw(channel, sigma):
    """单尺度SSR，不加epsilon"""
    channel = channel.astype(np.float32) / 255.0
    illumination = gaussian_filter(channel, sigma=sigma)
    with np.errstate(divide='ignore', invalid='ignore'):
        retinex = np.log(channel) - np.log(illumination)
        retinex[np.isneginf(retinex)] = np.min(retinex[np.isfinite(retinex)])
    return retinex

def msr_channel(channel, sigmas=[15, 80, 250]):
    """多尺度SSR"""
    retinex_scales = [ssr_channel_raw(channel, sigma) for sigma in sigmas]
    retinex = sum(retinex_scales) / len(sigmas)  # 等权平均
    # 归一化并转为 uint8
    retinex -= np.min(retinex)
    retinex /= np.max(retinex)
    return (retinex * 255).astype(np.uint8)

def msr_image(img_pil, sigmas=[15, 80, 250]):
    img = np.array(img_pil)
    if img.ndim == 2:
        return Image.fromarray(msr_channel(img, sigmas))
    elif img.ndim == 3:
        channels = [msr_channel(img[..., c], sigmas) for c in range(3)]
        retinex = np.stack(channels, axis=-1)
        return Image.fromarray(retinex)
    else:
        raise ValueError("Unsupported image format")

# 示例：使用 MSR 处理图像
if __name__ == "__main__":
    img = Image.open("D:\Retinex\\test\\test.jpg").convert("RGB")
    result = msr_image(img, sigmas=[15, 80, 250])

    # 可视化对比
    plt.subplot(1, 2, 1)
    plt.imshow(img)
    plt.title("Original")

    plt.subplot(1, 2, 2)
    plt.imshow(result)
    plt.title("MSR (3 scales)")
    plt.show()
